"""
数据分析模块

实现代表性数据分析和特征分析功能
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import StandardScaler
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


class DataAnalyzer:
    """数据分析器"""
    
    def __init__(self):
        self.representative_data = None
        self.analysis_results = {}
        
    def select_representative_data(self, source_data, n_samples_per_class=5):
        """
        选择代表性源域数据进行分析
        
        Args:
            source_data: 源域数据列表
            n_samples_per_class: 每个故障类型选择的样本数
        """
        print(f"🎯 选择代表性源域数据 (每类{n_samples_per_class}个样本)...")
        
        # 按故障类型分组
        data_by_type = {}
        for data in source_data:
            fault_type = data['fault_type']
            if fault_type not in data_by_type:
                data_by_type[fault_type] = []
            data_by_type[fault_type].append(data)
        
        # 从每个故障类型中选择代表性样本
        representative_data = []
        for fault_type, data_list in data_by_type.items():
            if len(data_list) >= n_samples_per_class:
                # 随机选择指定数量的样本
                selected = np.random.choice(data_list, n_samples_per_class, replace=False)
                representative_data.extend(selected)
                print(f"  {fault_type}: 从{len(data_list)}个样本中选择了{len(selected)}个")
            else:
                # 如果样本不足，选择全部
                representative_data.extend(data_list)
                print(f"  {fault_type}: 样本不足，选择了全部{len(data_list)}个")
        
        self.representative_data = representative_data
        print(f"✅ 共选择了{len(representative_data)}个代表性样本")
        
        return representative_data
    
    def analyze_representative_data(self, representative_data, feature_extractor):
        """
        分析代表性数据的特征
        
        Args:
            representative_data: 代表性数据列表
            feature_extractor: 特征提取器
        """
        print("📊 分析代表性数据特征...")
        
        # 提取特征
        features_list = []
        for data in representative_data:
            # 使用DE信号进行特征提取
            de_signal = data['de_signal']
            if len(de_signal) > 0:
                features = feature_extractor.extract_all_features(
                    de_signal, 
                    data['rpm'], 
                    'SKF6205'
                )
                features['file_name'] = data['file_name']
                features['fault_type'] = data['fault_type']
                features['fault_size'] = data['fault_size']
                features['load_condition'] = data['load_condition']
                features_list.append(features)
        
        # 转换为DataFrame
        features_df = pd.DataFrame(features_list)
        
        # 分析特征分布
        distribution_analysis = self._analyze_feature_distribution(features_df)
        
        # 分析特征相关性
        correlation_analysis = self._analyze_feature_correlation(features_df)
        
        # 分析故障类型特征差异
        fault_analysis = self._analyze_fault_characteristics(features_df)
        
        # 保存分析结果
        self.analysis_results = {
            'features_df': features_df,
            'distribution_analysis': distribution_analysis,
            'correlation_analysis': correlation_analysis,
            'fault_analysis': fault_analysis
        }
        
        return self.analysis_results
    
    def _analyze_feature_distribution(self, features_df):
        """分析特征分布"""
        print("  📈 分析特征分布...")
        
        # 获取数值特征列
        numeric_columns = features_df.select_dtypes(include=[np.number]).columns
        numeric_columns = [col for col in numeric_columns if col not in ['fault_size', 'load_condition']]
        
        # 按故障类型分析特征分布
        distribution_stats = {}
        for fault_type in features_df['fault_type'].unique():
            fault_data = features_df[features_df['fault_type'] == fault_type]
            fault_stats = {}
            
            for feature in numeric_columns:
                feature_data = fault_data[feature].dropna()
                if len(feature_data) > 0:
                    fault_stats[feature] = {
                        'mean': np.mean(feature_data),
                        'std': np.std(feature_data),
                        'min': np.min(feature_data),
                        'max': np.max(feature_data),
                        'median': np.median(feature_data),
                        'skewness': self._calculate_skewness(feature_data),
                        'kurtosis': self._calculate_kurtosis(feature_data)
                    }
            
            distribution_stats[fault_type] = fault_stats
        
        return distribution_stats
    
    def _analyze_feature_correlation(self, features_df):
        """分析特征相关性"""
        print("  📈 分析特征相关性...")
        
        # 获取数值特征列
        numeric_columns = features_df.select_dtypes(include=[np.number]).columns
        numeric_columns = [col for col in numeric_columns if col not in ['fault_size', 'load_condition']]
        
        # 计算相关系数矩阵
        corr_matrix = features_df[numeric_columns].corr()
        
        # 找出高相关性的特征对
        high_corr_pairs = []
        for i in range(len(corr_matrix.columns)):
            for j in range(i+1, len(corr_matrix.columns)):
                corr_val = corr_matrix.iloc[i, j]
                if abs(corr_val) > 0.8:  # 高相关性阈值
                    high_corr_pairs.append({
                        'feature1': corr_matrix.columns[i],
                        'feature2': corr_matrix.columns[j],
                        'correlation': corr_val
                    })
        
        return {
            'corr_matrix': corr_matrix,
            'high_corr_pairs': high_corr_pairs
        }
    
    def _analyze_fault_characteristics(self, features_df):
        """分析故障特征"""
        print("  📈 分析故障特征...")
        
        # 按故障类型分析特征
        fault_analysis = {}
        for fault_type in features_df['fault_type'].unique():
            fault_data = features_df[features_df['fault_type'] == fault_type]
            
            # 分析故障尺寸分布
            fault_sizes = fault_data['fault_size'].unique()
            load_conditions = fault_data['load_condition'].unique()
            
            fault_analysis[fault_type] = {
                'sample_count': len(fault_data),
                'fault_sizes': fault_sizes.tolist(),
                'load_conditions': load_conditions.tolist(),
                'file_names': fault_data['file_name'].tolist()
            }
        
        return fault_analysis
    
    def _calculate_skewness(self, data):
        """计算偏度"""
        mean = np.mean(data)
        std = np.std(data)
        if std == 0:
            return 0
        return np.mean(((data - mean) / std) ** 3)
    
    def _calculate_kurtosis(self, data):
        """计算峰度"""
        mean = np.mean(data)
        std = np.std(data)
        if std == 0:
            return 0
        return np.mean(((data - mean) / std) ** 4)
    
    def visualize_representative_analysis(self, analysis_results, save_dir='.'):
        """可视化代表性数据分析结果"""
        print("📊 生成代表性数据分析可视化...")
        
        features_df = analysis_results['features_df']
        distribution_analysis = analysis_results['distribution_analysis']
        correlation_analysis = analysis_results['correlation_analysis']
        
        # 1. 故障类型分布
        self._plot_fault_type_distribution(features_df, save_dir)
        
        # 2. 特征分布箱线图
        self._plot_feature_distribution_boxplot(features_df, save_dir)
        
        # 3. 特征相关性热力图
        self._plot_correlation_heatmap(correlation_analysis['corr_matrix'], save_dir)
        
        # 4. 特征空间可视化
        self._plot_feature_space_visualization(features_df, save_dir)
        
        # 5. 故障特征对比
        self._plot_fault_characteristics_comparison(features_df, save_dir)
    
    def _plot_fault_type_distribution(self, features_df, save_dir):
        """绘制故障类型分布图"""
        plt.figure(figsize=(10, 6))
        
        # 故障类型计数
        fault_counts = features_df['fault_type'].value_counts()
        
        # 创建柱状图
        bars = plt.bar(fault_counts.index, fault_counts.values, color=['skyblue', 'lightcoral', 'lightgreen', 'gold'])
        
        # 添加数值标签
        for bar, count in zip(bars, fault_counts.values):
            plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.1, 
                    str(count), ha='center', va='bottom', fontweight='bold')
        
        plt.title('代表性数据故障类型分布', fontsize=14, fontweight='bold')
        plt.xlabel('故障类型')
        plt.ylabel('样本数量')
        plt.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig(f'{save_dir}/representative_fault_distribution.png', dpi=300, bbox_inches='tight')
        plt.close()
        print("  ✅ 故障类型分布图已保存")
    
    def _plot_feature_distribution_boxplot(self, features_df, save_dir):
        """绘制特征分布箱线图"""
        # 选择前12个数值特征
        numeric_columns = features_df.select_dtypes(include=[np.number]).columns
        numeric_columns = [col for col in numeric_columns if col not in ['fault_size', 'load_condition']]
        selected_features = numeric_columns[:12]
        
        fig, axes = plt.subplots(3, 4, figsize=(16, 12))
        fig.suptitle('代表性数据特征分布箱线图', fontsize=16, fontweight='bold')
        
        axes = axes.flatten()
        
        for i, feature in enumerate(selected_features):
            if i < len(axes):
                # 按故障类型绘制箱线图
                data_by_type = []
                labels = []
                
                for fault_type in features_df['fault_type'].unique():
                    type_data = features_df[features_df['fault_type'] == fault_type][feature].dropna()
                    if len(type_data) > 0:
                        data_by_type.append(type_data)
                        labels.append(fault_type)
                
                if data_by_type:
                    axes[i].boxplot(data_by_type, labels=labels)
                    axes[i].set_title(f'{feature}')
                    axes[i].tick_params(axis='x', rotation=45)
                    axes[i].grid(True, alpha=0.3)
        
        # 隐藏多余的子图
        for i in range(len(selected_features), len(axes)):
            axes[i].set_visible(False)
        
        plt.tight_layout()
        plt.savefig(f'{save_dir}/representative_feature_distribution.png', dpi=300, bbox_inches='tight')
        plt.close()
        print("  ✅ 特征分布箱线图已保存")
    
    def _plot_correlation_heatmap(self, corr_matrix, save_dir):
        """绘制相关性热力图"""
        plt.figure(figsize=(12, 10))
        
        # 创建热力图
        mask = np.triu(np.ones_like(corr_matrix, dtype=bool))
        sns.heatmap(corr_matrix, mask=mask, annot=False, cmap='coolwarm', center=0,
                   square=True, linewidths=0.5, cbar_kws={"shrink": 0.8})
        
        plt.title('代表性数据特征相关性热力图', fontsize=14, fontweight='bold')
        plt.tight_layout()
        plt.savefig(f'{save_dir}/representative_correlation_heatmap.png', dpi=300, bbox_inches='tight')
        plt.close()
        print("  ✅ 特征相关性热力图已保存")
    
    def _plot_feature_space_visualization(self, features_df, save_dir):
        """绘制特征空间可视化"""
        # 获取数值特征
        numeric_columns = features_df.select_dtypes(include=[np.number]).columns
        numeric_columns = [col for col in numeric_columns if col not in ['fault_size', 'load_condition']]
        
        X = features_df[numeric_columns].values
        y = features_df['fault_type'].values
        
        # 标准化数据
        scaler = StandardScaler()
        X_scaled = scaler.fit_transform(X)
        
        # t-SNE降维
        tsne = TSNE(n_components=2, random_state=42, perplexity=min(30, len(X)-1))
        X_embedded = tsne.fit_transform(X_scaled)
        
        # 绘制散点图
        plt.figure(figsize=(12, 8))
        
        unique_labels = np.unique(y)
        colors = plt.cm.viridis(np.linspace(0, 1, len(unique_labels)))
        
        for i, label in enumerate(unique_labels):
            mask = y == label
            plt.scatter(X_embedded[mask, 0], X_embedded[mask, 1], 
                       c=[colors[i]], label=label, alpha=0.7, s=100)
        
        plt.xlabel('t-SNE 第一主成分')
        plt.ylabel('t-SNE 第二主成分')
        plt.title('代表性数据特征空间可视化 (t-SNE)', fontsize=14, fontweight='bold')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig(f'{save_dir}/representative_feature_space.png', dpi=300, bbox_inches='tight')
        plt.close()
        print("  ✅ 特征空间可视化图已保存")
    
    def _plot_fault_characteristics_comparison(self, features_df, save_dir):
        """绘制故障特征对比图"""
        # 选择几个关键特征进行对比
        key_features = ['mean', 'std', 'rms', 'kurtosis', 'freq_mean', 'approx_entropy']
        available_features = [f for f in key_features if f in features_df.columns]
        
        if not available_features:
            return
        
        fig, axes = plt.subplots(2, 3, figsize=(15, 10))
        fig.suptitle('代表性数据故障特征对比', fontsize=16, fontweight='bold')
        
        axes = axes.flatten()
        
        for i, feature in enumerate(available_features[:6]):
            if i < len(axes):
                # 按故障类型绘制特征对比
                fault_types = features_df['fault_type'].unique()
                feature_data = []
                labels = []
                
                for fault_type in fault_types:
                    type_data = features_df[features_df['fault_type'] == fault_type][feature].dropna()
                    if len(type_data) > 0:
                        feature_data.append(type_data)
                        labels.append(fault_type)
                
                if feature_data:
                    axes[i].boxplot(feature_data, labels=labels)
                    axes[i].set_title(f'{feature}')
                    axes[i].tick_params(axis='x', rotation=45)
                    axes[i].grid(True, alpha=0.3)
        
        # 隐藏多余的子图
        for i in range(len(available_features), len(axes)):
            axes[i].set_visible(False)
        
        plt.tight_layout()
        plt.savefig(f'{save_dir}/representative_fault_characteristics.png', dpi=300, bbox_inches='tight')
        plt.close()
        print("  ✅ 故障特征对比图已保存")
    
    def generate_representative_analysis_report(self, analysis_results, save_path=None):
        """生成代表性数据分析报告"""
        print("📝 生成代表性数据分析报告...")
        
        features_df = analysis_results['features_df']
        distribution_analysis = analysis_results['distribution_analysis']
        correlation_analysis = analysis_results['correlation_analysis']
        fault_analysis = analysis_results['fault_analysis']
        
        report = f"""# 代表性源域数据分析报告

## 1. 数据概况

### 1.1 样本统计
- **总样本数**: {len(features_df)}
- **故障类型数**: {len(features_df['fault_type'].unique())}
- **特征维度**: {len(features_df.select_dtypes(include=[np.number]).columns) - 2}

### 1.2 故障类型分布
"""
        
        # 添加故障类型分布
        fault_counts = features_df['fault_type'].value_counts()
        for fault_type, count in fault_counts.items():
            percentage = count / len(features_df) * 100
            report += f"- **{fault_type}**: {count}个样本 ({percentage:.1f}%)\n"
        
        report += f"""
## 2. 特征分布分析

### 2.1 各故障类型特征统计
"""
        
        # 添加特征分布统计
        for fault_type, stats in distribution_analysis.items():
            report += f"\n**{fault_type}**:\n"
            # 选择几个关键特征
            key_features = ['mean', 'std', 'rms', 'kurtosis']
            for feature in key_features:
                if feature in stats:
                    feature_stats = stats[feature]
                    report += f"  - {feature}: 均值={feature_stats['mean']:.4f}, 标准差={feature_stats['std']:.4f}\n"
        
        report += f"""
## 3. 特征相关性分析

### 3.1 高相关性特征对
发现 {len(correlation_analysis['high_corr_pairs'])} 对高相关性特征 (|相关系数| > 0.8):

"""
        
        # 添加高相关性特征对
        for i, pair in enumerate(correlation_analysis['high_corr_pairs'][:10], 1):
            report += f"{i}. {pair['feature1']} ↔ {pair['feature2']}: {pair['correlation']:.4f}\n"
        
        if len(correlation_analysis['high_corr_pairs']) > 10:
            report += f"... 还有 {len(correlation_analysis['high_corr_pairs']) - 10} 对高相关性特征\n"
        
        report += f"""
## 4. 故障特征分析

### 4.1 各故障类型特征
"""
        
        # 添加故障特征分析
        for fault_type, analysis in fault_analysis.items():
            report += f"\n**{fault_type}**:\n"
            report += f"  - 样本数量: {analysis['sample_count']}\n"
            report += f"  - 故障尺寸: {analysis['fault_sizes']}\n"
            report += f"  - 载荷条件: {analysis['load_conditions']}\n"
            report += f"  - 文件列表: {', '.join(analysis['file_names'])}\n"
        
        report += f"""
## 5. 分析结论

### 5.1 主要发现
1. **数据代表性**: 成功选择了{len(features_df)}个代表性样本，覆盖了所有故障类型
2. **特征分布**: 不同故障类型在关键特征上表现出明显差异
3. **特征相关性**: 发现{len(correlation_analysis['high_corr_pairs'])}对高相关性特征，需要关注特征冗余问题
4. **故障特征**: 各故障类型具有独特的特征模式，为后续分类提供了基础

### 5.2 后续建议
1. **特征选择**: 基于代表性数据分析结果，选择最具区分性的特征
2. **数据增强**: 可以考虑对样本较少的故障类型进行数据增强
3. **特征工程**: 可以基于故障机理进一步优化特征提取方法
4. **模型设计**: 考虑不同故障类型的特征差异，设计针对性的分类模型

---
*报告生成时间: {pd.Timestamp.now().strftime('%Y年%m月%d日 %H:%M:%S')}*
"""
        
        if save_path:
            with open(save_path, 'w', encoding='utf-8') as f:
                f.write(report)
            print(f"✅ 代表性数据分析报告已保存: {save_path}")
        
        return report


